#!/usr/bin/env python3
"""
部署环境检查脚本
将项目上传到云服务器上之后运行这个脚本检测该项目在云服务器中的环境是否正常
正常之后就可以部署项目
"""

import os
import sys
from pathlib import Path
import importlib
import subprocess

def get_project_root():
    """获取项目根目录路径"""
    current_file = Path(__file__).resolve()
    return current_file.parent

def check_python_version():
    """检查Python版本"""
    print("=" * 50)
    print("1. 检查Python版本")
    print("=" * 50)
    version = sys.version_info
    print(f"当前Python版本: {version.major}.{version.minor}.{version.micro}")
    if version.major == 3 and version.minor >= 8:
        print("✓ Python版本符合要求 (>= 3.8)")
        return True
    else:
        print("✗ Python版本过低，需要 >= 3.8")
        return False

def check_required_packages():
    """检查必需的Python包"""
    print("\n" + "=" * 50)
    print("2. 检查Python依赖包")
    print("=" * 50)
    
    required_packages = [
        'fastapi', 'uvicorn', 'pydantic', 'requests',
        'torch', 'torchaudio', 'transformers', 'sentence_transformers',
        'librosa', 'opensmile', 'audonnx', 'faster_whisper',
        'numpy', 'pandas', 'scikit_learn', 'joblib',
        'opencc', 'jieba', 'pydub'
    ]
    
    missing_packages = []
    for package in required_packages:
        try:
            # 特殊处理一些包名
            if package == 'scikit_learn':
                importlib.import_module('sklearn')
            elif package == 'faster_whisper':
                importlib.import_module('faster_whisper')
            else:
                importlib.import_module(package)
            print(f"✓ {package}")
        except ImportError:
            print(f"✗ {package} - 未安装")
            missing_packages.append(package)
    
    if missing_packages:
        print(f"\n缺少以下包，请安装: {', '.join(missing_packages)}")
        print("运行: pip install -r requirements.txt")
        return False
    else:
        print("\n✓ 所有Python依赖包已安装")
        return True

def check_system_dependencies():
    """检查系统依赖"""
    print("\n" + "=" * 50)
    print("3. 检查系统依赖")
    print("=" * 50)
    
    dependencies = {
        'ffmpeg': 'ffmpeg -version',
    }
    
    all_ok = True
    for dep, cmd in dependencies.items():
        try:
            result = subprocess.run(cmd.split(), capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                print(f"✓ {dep} - 已安装")
            else:
                print(f"✗ {dep} - 未正确安装")
                all_ok = False
        except (subprocess.TimeoutExpired, FileNotFoundError):
            print(f"✗ {dep} - 未安装或不在PATH中")
            all_ok = False
    
    # 检查OpenSMILE
    try:
        import opensmile
        smile = opensmile.Smile(
            feature_set=opensmile.FeatureSet.eGeMAPSv02,
            feature_level=opensmile.FeatureLevel.Functionals,
        )
        print("✓ OpenSMILE - 可正常使用")
    except Exception as e:
        print(f"✗ OpenSMILE - 初始化失败: {str(e)}")
        all_ok = False
    
    return all_ok

def check_model_files():
    """检查模型文件"""
    print("\n" + "=" * 50)
    print("4. 检查模型文件")
    print("=" * 50)
    
    project_root = get_project_root()
    model_dir = project_root / "model"
    
    required_models = {
        'w2v2-L-robust': '情感分析模型',
        'faster-whisper-medium': 'ASR语音识别模型',
        'scoring_model.pth': '多模态评分模型',
        'tech_stack_model': '技术栈提取模型',
        'sbert_all_MiniLM_L6_v2': 'SBERT文本嵌入模型'
    }
    
    all_models_ok = True
    for model_name, description in required_models.items():
        model_path = model_dir / model_name
        if model_path.exists():
            print(f"✓ {model_name} - {description}")
        else:
            print(f"✗ {model_name} - {description} (未找到)")
            all_models_ok = False
    
    return all_models_ok

def check_network_connectivity():
    """检查网络连接"""
    print("\n" + "=" * 50)
    print("5. 检查网络连接")
    print("=" * 50)
    
    import requests
    
    test_urls = [
        "https://www.baidu.com",
        "https://huggingface.co",
    ]
    
    all_ok = True
    for url in test_urls:
        try:
            response = requests.get(url, timeout=10)
            if response.status_code == 200:
                print(f"✓ {url} - 可访问")
            else:
                print(f"✗ {url} - 响应码: {response.status_code}")
                all_ok = False
        except Exception as e:
            print(f"✗ {url} - 连接失败: {str(e)}")
            all_ok = False
    
    return all_ok

def test_model_loading():
    """测试模型加载"""
    print("\n" + "=" * 50)
    print("6. 测试模型加载")
    print("=" * 50)
    
    project_root = get_project_root()
    sys.path.insert(0, str(project_root))
    
    try:
        # 测试情感分析模型
        from src.service.analyze_audio_emotion import RobustAudioEmotionAnalyzer, get_model_path
        try:
            analyzer = RobustAudioEmotionAnalyzer(
                model_path=get_model_path("w2v2-L-robust"),
                window_size=3,
                chunk_size=1.8
            )
            print("✓ 情感分析模型加载成功")
        except Exception as e:
            print(f"✗ 情感分析模型加载失败: {str(e)}")
            return False
        
        # 测试压力检测模型
        from src.service.stress_detector import StressDetector
        try:
            stress_detector = StressDetector()
            print("✓ 压力检测模型加载成功")
        except Exception as e:
            print(f"✗ 压力检测模型加载失败: {str(e)}")
            return False
        
        # 测试ASR模型
        from src.service.whisper_asr import WhisperASR
        try:
            asr = WhisperASR(model_path=get_model_path("faster-whisper-medium"))
            print("✓ ASR语音识别模型加载成功")
        except Exception as e:
            print(f"✗ ASR语音识别模型加载失败: {str(e)}")
            return False
        
        # 测试多模态评分模型
        from src.service.multimodal.predictor import InterviewPredictor
        try:
            predictor = InterviewPredictor(model_path=get_model_path("scoring_model.pth"))
            print("✓ 多模态评分模型加载成功")
        except Exception as e:
            print(f"✗ 多模态评分模型加载失败: {str(e)}")
            return False
        
        # 测试技术栈提取模型
        from src.service.extract_tech_stack.predictor import TechStackPredictor, jieba_tokenizer
        try:
            # 解决pickle加载问题：将jieba_tokenizer函数添加到当前__main__模块
            sys.modules['__main__'].jieba_tokenizer = jieba_tokenizer
            
            tech_predictor = TechStackPredictor(model_dir=get_model_path("tech_stack_model"))
            print("✓ 技术栈提取模型加载成功")
        except Exception as e:
            print(f"✗ 技术栈提取模型加载失败: {str(e)}")
            return False
        
        print("\n✓ 所有模型加载测试通过")
        return True
        
    except Exception as e:
        print(f"✗ 模型加载测试失败: {str(e)}")
        return False

def main():
    """主函数"""
    print("AI面试后端项目 - 部署环境检查")
    print("检查时间:", __import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
    
    checks = [
        ("Python版本", check_python_version),
        ("Python依赖包", check_required_packages),
        ("系统依赖", check_system_dependencies),
        ("模型文件", check_model_files),
        ("网络连接", check_network_connectivity),
        ("模型加载", test_model_loading),
    ]
    
    results = {}
    for check_name, check_func in checks:
        try:
            results[check_name] = check_func()
        except Exception as e:
            print(f"✗ {check_name} 检查过程中出错: {str(e)}")
            results[check_name] = False
    
    # 总结
    print("\n" + "=" * 50)
    print("检查结果总结")
    print("=" * 50)
    
    all_passed = True
    for check_name, passed in results.items():
        status = "✓ 通过" if passed else "✗ 失败"
        print(f"{check_name:15s}: {status}")
        if not passed:
            all_passed = False
    
    print("\n" + "=" * 50)
    if all_passed:
        print("🎉 所有检查通过！项目可以正常部署运行。")
        print("启动服务器: python src/service/server_api.py")
    else:
        print("❌ 部分检查失败，请根据上述提示解决问题后重新检查。")
    print("=" * 50)
    
    return all_passed

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1) 